#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
批量：FITS 星表 -> 查询光变曲线（调用现有 LightCurveFetcher）
- FITS 列名默认: Name, RAdeg, DEdeg（单位：度）
- 将 RA/Dec 转为 'hh:mm:ss.ssssss' / '±dd:mm:ss.ssssss' 传入接口
- 支持只处理前 N 条（limit）、从第 N 行开始（skip）、失败不终止、弹窗画图与保存 PNG
- 汇总结果保存为 parquet

依赖：
- loguru, pandas, matplotlib, astropy, tqdm
- 你已有的 main_light_curve.py 中的 LightCurveFetcher / plot_light_curve / plot_multiband
"""

import os
import time
from typing import List, Dict, Optional

from loguru import logger
from matplotlib import pyplot as plt
import pandas as pd
from astropy.table import Table
from astropy.coordinates import SkyCoord
import astropy.units as u
from tqdm import tqdm

# 你已有的实现（保持不变）
from main_light_curve import LightCurveFetcher, plot_light_curve, plot_multiband


# ========================= FITS 读取 & 工具 ========================= #

def load_targets_from_fits(
    fits_path: str,
    name_col: str = "Name",
    ra_col: str = "RAdeg",
    dec_col: str = "DEdeg",
    limit: Optional[int] = None,
    skip: int = 0,  # NEW: 跳过前 N 行
) -> List[Dict]:
    """
    读取 FITS 星表，返回 [{name, ra_hms, dec_dms, ra_deg, dec_deg}] 列表。
    :param fits_path: FITS 星表路径
    :param limit: 只取接下来多少行（None 表示到文件尾）
    :param skip: 跳过前 N 行（从第 skip 行开始处理）
    """
    tab = Table.read(fits_path, format="fits")
    total_len = len(tab)

    if skip >= total_len:
        raise ValueError(f"skip={skip} 已超过星表总行数 {total_len}")

    start_idx = skip
    end_idx = total_len if limit is None else min(total_len, skip + limit)

    rows = []
    for i in range(start_idx, end_idx):
        name = str(tab[name_col][i]) if name_col in tab.colnames else f"row_{i}"
        ra_deg = float(tab[ra_col][i])
        dec_deg = float(tab[dec_col][i])
        c = SkyCoord(ra_deg * u.deg, dec_deg * u.deg, frame="icrs")
        ra_hms = c.ra.to_string(unit=u.hour, sep=":", precision=6, pad=True)
        dec_dms = c.dec.to_string(unit=u.deg, sep=":", precision=6, pad=True, alwayssign=True)
        rows.append({"name": name, "ra": ra_hms, "dec": dec_dms, "ra_deg": ra_deg, "dec_deg": dec_deg})

    logger.info(f"从第 {start_idx} 行开始，读取到第 {end_idx-1} 行，共 {len(rows)} 个目标。")
    return rows


def _ensure_errcol(df: pd.DataFrame) -> pd.DataFrame:
    """确保存在误差列 magerr_auto_s；若没有则回退或填充空值。"""
    if "magerr_auto_s" in df.columns:
        return df
    for alt in ["magerr", "mag_err", "MAGERR", "magerr_auto"]:
        if alt in df.columns:
            return df.rename(columns={alt: "magerr_auto_s"})
    df = df.copy()
    df["magerr_auto_s"] = pd.NA
    return df


# ========================= 批量主流程 ========================= #

def batch_query_from_fits(
    fits_path: str,
    base_url: str,
    radius: float = 1.5,
    band: str = "all",
    start_date: Optional[str] = None,
    save_dir: str = "./lightcurves",
    limit: Optional[int] = None,
    skip: int = 0,             # NEW: 跳过行数
    pause_s: float = 0.25,
    show_plot: bool = True,    # 是否弹窗
    save_plot: bool = True,    # 是否保存 PNG
) -> pd.DataFrame:
    """
    批量查询：
    - 每个目标查询成功后即时画图（band='all' -> 多波段；否则只画该 band）
    - 保存 PNG 至 save_dir/figs/<Name>.png
    - 汇总所有记录到 save_dir/all_targets_lightcurves.parquet
    """
    os.makedirs(save_dir, exist_ok=True)
    figdir = os.path.join(save_dir, "figs")
    os.makedirs(figdir, exist_ok=True)

    if show_plot:
        plt.ion()

    fetcher = LightCurveFetcher(base_url)
    targets = load_targets_from_fits(fits_path, limit=limit, skip=skip)
    all_rows = []

    for t in tqdm(targets, desc="Querying light curves"):
        try:
            df = fetcher.fetch(
                t["ra"], t["dec"],
                radius=radius, band=band, start_date=start_date, download=False
            )
            if not isinstance(df, pd.DataFrame) or df.empty:
                logger.warning(f"[WARN] {t['name']} 无数据")
                time.sleep(pause_s)
                continue

            df = _ensure_errcol(df)
            # 附加标识列
            df = df.copy()
            df["target_name"] = t["name"]
            df["RA"] = t["ra"]; df["Dec"] = t["dec"]
            df["RAdeg"] = t["ra_deg"]; df["DEdeg"] = t["dec_deg"]
            all_rows.append(df)

            # 画图（使用你已有的函数）
            title = f"{t['name']}  (RA={t['ra']}, Dec={t['dec']})"
            if band == "all" and ("band" in df.columns):
                plot_multiband(df)
            else:
                plot_light_curve(df, title=title)

            # 保存 PNG
            if save_plot:
                safe_name = t["name"].replace(" ", "_").replace("/", "_")
                plt.gcf().savefig(os.path.join(figdir, f"{safe_name}.png"), dpi=130)

            if show_plot:
                plt.pause(0.001)
            plt.close()

        except Exception as e:
            logger.error(f"[WARN] {t['name']} 查询失败：{e}")

        time.sleep(pause_s)

    # 汇总保存
    if all_rows:
        df_all = pd.concat(all_rows, ignore_index=True)
        out_parq = os.path.join(save_dir, "all_targets_lightcurves.parquet")
        df_all.to_parquet(out_parq, index=False)
        logger.info(f"✅ 汇总保存：{out_parq}，总记录 {len(df_all)}")
        return df_all
    else:
        logger.warning("⚠️ 没有获得任何记录")
        return pd.DataFrame()


# ========================= 示例入口 ========================= #

if __name__ == "__main__":
    # ✅ 基本配置
    BASE_URL   = "http://192.168.16.70:8082/api/query_light_curve_xp_version_V20250303_KM"
    RADIUS     = 1.5        # arcsec
    BAND       = "all"      # 'r', 'g', 'all'
    START_DATE = None       # or "2025-07-01"

    # ✅ 批量配置
    FITS_PATH  = "AllVariable.fits"   # 改为你的星表路径
    LIMIT      = 100000                  # 本轮处理多少行（None 表示直到文件尾）
    SKIP       = 2000+20000+100000                  # 跳过前 200 行（从第 201 行开始）
    SAVE_DIR   = "./lightcurves"

    df_all = batch_query_from_fits(
        fits_path=FITS_PATH,
        base_url=BASE_URL,
        radius=RADIUS,
        band=BAND,
        start_date=START_DATE,
        save_dir=SAVE_DIR,
        limit=LIMIT,
        skip=SKIP,           # << NEW
        pause_s=0.2,
        show_plot=True,      # 大量目标建议 False（只保存 PNG）
        save_plot=True,
    )
    logger.info(f"✅ 批量完成，合并记录数：{len(df_all)}")
